# @file name  : decorator_demo.py
# @brief      : 对图片进行分割推理
"""

# 导入必要的库
from flask import Flask, render_template, request  
import time  
import os  
import cv2  
import torch  
from tools.predictor import Predictor  # 导入自定义的神经网络模型

# 初始化
BASEDIR = os.path.dirname(__file__)  # 获取当前脚本所在目录
app = Flask(__name__)  # 创建一个Flask应用
upload_dir = os.path.join(BASEDIR, "static", "upload_img")  # 定义上传图片的目录
if not os.path.exists(upload_dir):  # 如果目录不存在，则创建它
    os.makedirs(upload_dir)

path_checkpoint = r"G:\img_seg\seg_baseline\checkpoint_best.pkl"  # 预训练模型的路径

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # 根据GPU是否可用选择设备
predictor = Predictor(path_checkpoint, device=device)  # 创建图像分割模型的实例
in_size = 512  # 输入图像大小

# 以下是一些辅助函数，用于保存图像、生成HTML页面等，这些函数没有涉及到核心逻辑，故不一一注释

# 定义该url接收get和post请求，可到route的add_url_rule函数中看到默认是get请求
@app.route("/", methods=["GET", "POST"])  # 定义根路径的路由，接受GET和POST请求
def func():
    if request.method == "POST":  # 如果请求方法是POST
        try:
            f = request.files['imgfile']  # 从请求中获取上传的图像文件
            path_img, file_name = save_img(f, upload_dir)  # 保存上传的图像文件
            img_t, img_bgr = predictor.preprocess(path_img, in_size=in_size, keep_ratio=True)  # 预处理图像
            _, pred_mask = predictor.predict(img_t)  # 使用模型进行图像分割
            img_matting = predictor.postprocess(img_bgr, pred_mask, color="w")  # 后处理分割结果
            name_html = gen_html(img_matting, img_bgr, file_name, upload_dir)  # 生成用于展示的HTML页面
            return render_template(name_html)  # 返回HTML页面
        except Exception as e:
            return f"{e}, Please try it again!"  # 发生异常时返回错误信息
    else:
        return render_template("upload.html")  # 如果请求方法是GET，返回一个上传图片的表单页面

if __name__ == '__main__':
    app.run()  # 启动Flask应用
    # 允许外部访问，但无公网IP，仅局域网内其他主机可访问，如同wifi下的设备，本机IP可通过ipconfig命令查看，mac通过ipconfig /a
    # app.run(host="0.0.0.0", port=80)  # 允许公网访问的示例，端口80